﻿/*
题目：寻找两个正序数组的中位数

给定两个大小分别为 m 和 n 的正序（从小到大）数组 nums1 和 nums2。请你找出并返回这两个正序数组的 中位数 。

算法的时间复杂度应该为 O(log (m+n)) 。

https://leetcode.cn/problems/median-of-two-sorted-arrays/description/
*/

#include <iostream>
#include <string>
#include <vector>
#include <list>
#include <array>
#include "TreeNode.hpp"
#include "ListNode.hpp"
#include <map>
#include <set>
#include <stack>
#include <queue>
#include <unordered_map>
#include <unordered_set>
#include <numeric>
#include <algorithm>
#include <sstream>
#include <functional>
#include <random>

using namespace std;

class Solution {
public:
    // 相当于在两个有序数组中找 第K大的数字，K=中位数
    double findMedianSortedArrays(vector<int>& nums1, vector<int>& nums2) {
        int total_len = nums1.size() + nums2.size();

        pair<int, int> mid;
        if ((total_len & 1) == 1) {
            mid.first = mid.second = total_len / 2 + 1;
        }
        else {
            mid.first = total_len / 2;
            mid.second = total_len / 2 + 1;
        }

        auto test1 = findKth(nums1, 0, nums2, 0, mid.first);
        auto test2 = findKth(nums1, 0, nums2, 0, mid.second);

        return (double)(test1 + test2) / 2;
    }

    // 从两个有序数组中找到第 K 大的元素
    int findKth(const vector<int>& left, int lpos, const vector<int>& right, int rpos, int k) {
        // 如果已经超出了左边的搜索范围
        if (lpos >= left.size()) {
            return right[rpos + k - 1];
        }
        if (rpos >= right.size()) {
            return left[lpos + k - 1];
        }

        // 如果 K 到了最小值
        if (k == 1) {
            return min(left[lpos], right[rpos]);
        }

        // 开始淘汰 前 K/2 个元素, 通过比较各自的 K/2 大的元素
        // 如果左边的元素不够比，那么就可以把右边 K/2 个元素排掉
        int left_num = lpos + k / 2 - 1 >= left.size() ? INT_MAX : left[lpos + k / 2 - 1];
        int right_num = rpos + k / 2 - 1 >= right.size() ? INT_MAX : right[rpos + k / 2 - 1];

        return left_num > right_num ?
            findKth(left, lpos, right, rpos + k / 2, k - k / 2) :
            findKth(left, lpos + k / 2, right, rpos, k - k / 2);
    }
};
